import os
import builtins
import sys
import warnings
import random
import json

# iter all dataset
def iter_dataset(dataset_dir: str):
    parts=[os.path.join(dataset_dir,e) for e in  os.listdir(dataset_dir)]
    
    result=[]
    
    for part in parts:
        good_dir = os.path.join(part, 'good')
        bad_dir = os.path.join(part, 'bad')
        
        good_examples = [os.path.join(good_dir, e) for e in os.listdir(good_dir)]
        bad_examples = [os.path.join(bad_dir, e) for e in os.listdir(bad_dir)]
        
        for example in good_examples:
            cdg_graph = os.path.join(example, 'cdg_graph.json')
            cfg_graph = os.path.join(example, 'cfg_graph.json')
            cg_graph = os.path.join(example, 'cg_graph.json')
            ddg_graph = os.path.join(example, 'ddg_graph.json')
            
            label=0
            
            result.append(([cdg_graph, cfg_graph, cg_graph, ddg_graph],label))
            
        for example in bad_examples:
            cdg_graph = os.path.join(example, 'cdg_graph.json')
            cfg_graph = os.path.join(example, 'cfg_graph.json')
            cg_graph = os.path.join(example, 'cg_graph.json')
            ddg_graph = os.path.join(example, 'ddg_graph.json')
            
            label=1
            
            result.append(([cdg_graph, cfg_graph, cg_graph, ddg_graph],label))
            
    return result

# iter all notation
def iter_notation(notation_dir: str):
    pass

# random get samples for visualization
def random_get_samples(dataset_dir: str, each_sample_num=3, sample_file_list=r'preprocessed_data/vis/sample_file_list.json'):
    
    if os.path.exists(sample_file_list):
        f = open(sample_file_list)
        file_list = json.load(f)
        return file_list

    parts=[os.path.join(dataset_dir,e) for e in os.listdir(dataset_dir)]
    # NOTE CAN BE BETTER:   ONLY ONE s0* INCLUDED
    random.shuffle(parts)
    result=[]
    good_num, bad_num = 0, 0
    
    for part in parts:
        good_dir = os.path.join(part, 'good')
        bad_dir = os.path.join(part, 'bad')
        
        good_examples = [os.path.join(good_dir, e) for e in os.listdir(good_dir)]
        bad_examples = [os.path.join(bad_dir, e) for e in os.listdir(bad_dir)]
        random.shuffle(good_examples)
        random.shuffle(bad_examples)
        
        for example in good_examples:
            cdg_graph = os.path.join(example, 'cdg_graph.json')
            cfg_graph = os.path.join(example, 'cfg_graph.json')
            cg_graph = os.path.join(example, 'cg_graph.json')
            ddg_graph = os.path.join(example, 'ddg_graph.json')
            
            label=0
            
            if good_num < each_sample_num:
                result.append(([cdg_graph, cfg_graph, cg_graph, ddg_graph],label))
                good_num += 1
            else:
                break
            
            
        for example in bad_examples:
            cdg_graph = os.path.join(example, 'cdg_graph.json')
            cfg_graph = os.path.join(example, 'cfg_graph.json')
            cg_graph = os.path.join(example, 'cg_graph.json')
            ddg_graph = os.path.join(example, 'ddg_graph.json')
            
            label=1
            
            if bad_num < each_sample_num:
                result.append(([cdg_graph, cfg_graph, cg_graph, ddg_graph],label))
                bad_num += 1
            else:
                break
            
    f = open(sample_file_list, 'w')
    json.dump(result, f)
    return result


def get_all_samples(dataset_dir:str, sample_file_list:str):
    result = iter_dataset(dataset_dir)
    f = open(sample_file_list, 'w')
    json.dump(result, f)
    return result

# ez modified print
def init_log(log_dir=r'logs/run_log.txt'):
    original_print = builtins.print
    def custom_print(*args, **kwargs):
        original_print(*args, **kwargs)
        with open(log_dir, 'a') as log_file:
            log_file.write(' '.join(map(str, args)) + '\n')
    builtins.print = custom_print
    def exception_handler(exctype, value, traceback):
        print(f"ERROR: {value}\n")
    sys.excepthook = exception_handler
    def warning_handler(message, category, filename, lineno, file=None, line=None):
        print(f"WARNING: {message}\n")
    warnings.showwarning = warning_handler